import os
import json
import copy
import torch

import logging


logger = logging.getLogger(__file__)


class ModelConfig(object):
    def __init__(self, **kwargs):
        pass

    def _load_finish_(self):
        pass

    @classmethod
    def from_dict(cls, json_object):
        config = cls()
        for key, value in json_object.items():
            if hasattr(config, key):
                config.__dict__[key] = value
        config._load_finish_()
        return config

    @classmethod
    def from_json_file(cls, json_file):
        with open(json_file, "r", encoding="utf-8") as reader:
            text = reader.read()
        return cls.from_dict(json.loads(text))

    def to_dict(self):
        output = copy.deepcopy(self.__dict__)
        output = {k: v for k, v in output.items() if v is not None}
        return output

    def to_json_string(self):
        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

    def to_json_file(self, json_file_path):
        with open(json_file_path, "w", encoding='utf-8') as writer:
            writer.write(self.to_json_string())


class ModelBase(torch.nn.Module):
    config_class = ModelConfig
    base_model_prefix = ''

    def __init__(self, config: ModelConfig, **kwargs):
        super(ModelBase, self).__init__()
        self.config = config

    def config_json_string(self):
        return self.config.to_json_string()

    def save_pretrained(self, model_dir, model_name='pytorch_model.bin', config_name='config.json'):
        """ 保存训练完成的模型 """
        if not os.path.exists(model_dir):
            raise ValueError(f"模型存储文件夹不存在，请注意:{model_dir}")
        config_path = os.path.join(model_dir, config_name)
        model_path = os.path.join(model_dir, model_name)
        torch.save(self.state_dict(), model_path)
        self.config.to_json_file(config_path)
        logger.info(f"save pretrained model to {model_dir}")

    @classmethod
    def from_pretrained(cls, model_dir, model_name='pytorch_model.bin', config_name='config.json', **kwargs):
        """ 简要处理从预训练模型中获取 """
        config_path = os.path.join(model_dir, config_name)
        model_path = os.path.join(model_dir, model_name)

        if not os.path.exists(config_path):
            raise ValueError(f"模型配置文件不存在，请检查: {config_path}")
        config = cls.config_class.from_json_file(config_path)

        for k, v in kwargs.items():
            if hasattr(config, k):
                config.__dict__[k] = v

        model = cls(config=config)

        if not os.path.exists(model_path):
            print(f"模型文件不存在，将不加载模型权值")
            return model

        # 处理前缀
        prefix = cls.base_model_prefix
        if (len(prefix) > 0) and (not str(prefix).endswith('.')):
            prefix += "."

        state_dict = torch.load(model_path)
        new_state_dict = {}
        for k, v in state_dict.items():
            if 'gamma' in k:
                k = k.replace('gamma', 'weight')
            if 'beta' in k:
                k = k.replace('beta', 'bias')
            if len(prefix) > 0:
                # key 开始能匹配到
                if str(k).startswith(prefix):
                    # 模型无此参数
                    if not hasattr(model, prefix.split('.')[0]):
                        k = k[len(prefix):]
            new_state_dict[k] = v
        info = model.load_state_dict(new_state_dict, strict=False)
        if info:
            print(f"load model for pretrained: load state info: \n\t{info}")
            logger.info(info)

        return model
